#include "detector/detector_node.hpp"

#include <string>
#include <vector>

#include <cv_bridge/cv_bridge.h>
#include <Eigen/Core>

#include "myutils/utils.hpp"
#include "rm_interfaces/AllTargets.h"
#include "rm_interfaces/Armor.h"

namespace detector {
DetectorNode::DetectorNode() : nh_("~"), enemy_color_(0) {
  nh_.param("/debug", debug_, true);

  std::string model_path;
  nh_.getParam("/detector/model_path", model_path);

  nh_.param("/detector/min_lightness", lightness_thresh_, 120);
  double a_max_angle, min_small_center_distance, max_small_center_distance,
      min_large_center_distance, max_large_center_distance, min_light_ratio;
  nh_.param("/detector/armor/max_angle", a_max_angle, 40.0);
  nh_.param("/detector/armor/min_small_center_distance",
            min_small_center_distance, 0.6);
  nh_.param("/detector/armor/max_small_center_distance",
            max_small_center_distance, 4.0);
  nh_.param("/detector/armor/min_large_center_distance",
            min_large_center_distance, 4.0);
  nh_.param("/detector/armor/max_large_center_distance",
            max_large_center_distance, 8.0);
  nh_.param("/detector/armor/min_light_ratio", min_light_ratio, 0.5);

  double l_max_angle, min_ratio, max_ratio;
  nh_.param("/detector/light/max_angle", l_max_angle, 40.0);
  nh_.param("/detector/min_ratio", min_ratio, 0.1);
  nh_.param("/detector/max_ratio", max_ratio, 1.0);

  Detector::LightParams l_params{
      .min_ratio = min_ratio, .max_ratio = max_ratio, .max_angle = l_max_angle};
  Detector::ArmorParams a_params{
      .min_light_ratio = min_light_ratio,
      .min_small_center_distance = min_small_center_distance,
      .max_small_center_distance = max_small_center_distance,
      .min_large_center_distance = min_large_center_distance,
      .max_large_center_distance = max_large_center_distance,
      .max_angle = a_max_angle};
  detector_ = std::make_unique<Detector>(lightness_thresh_, l_params, a_params,
                                         model_path);

  set_mode_server_ = nh_.advertiseService("/detector/set_mode",
                                          &DetectorNode::set_mode_cb, this);

  img_sub_ = nh_.subscribe<sensor_msgs::Image>("/raw_img", 10,
                                               &DetectorNode::image_cb, this);

  targets_pub_ = nh_.advertise<rm_interfaces::AllTargets>("/all_targets", 10);

  f_ = boost::bind(&DetectorNode::configure_cb, this, _1, _2);
  config_server_.setCallback(f_);

  if (debug_) {
    binary_img_pub_ = nh_.advertise<sensor_msgs::Image>("/binary_img", 1);
    rgb_img_pub_ = nh_.advertise<sensor_msgs::Image>("/detect_result", 1);
    number_img_pub_ = nh_.advertise<sensor_msgs::Image>("/number_img", 1);
  }
  last_time_ = ros::Time::now();
}

bool DetectorNode::set_mode_cb(rm_interfaces::SetMode::Request &request,
                               rm_interfaces::SetMode::Response &response) {
  switch (request.mode) {
    case 0: {
      autoaim_start_ = true;
      enemy_color_ = RED;
      response.success = true;
    } break;
    case 1: {
      autoaim_start_ = true;
      enemy_color_ = BLUE;
      response.success = true;
    } break;
    case 2:
    case 3:
    case 4:
    case 5: {
      autoaim_start_ = false;
      response.success = true;
    } break;
    default:
      response.success = false;
  }
  return true;
}

void DetectorNode::image_cb(const sensor_msgs::ImageConstPtr &msg) {
  ros::Time start_time = ros::Time::now();

  if (!autoaim_start_) {
    return;
  }
  // get src
  cv::Mat src = cv_bridge::toCvShare(msg, "bgr8")->image;

  // detect armors
  cv::Mat binary_img =
      detector_->preprocessImage(src, (enemy_color_ == RED ? 2 : 0));
  std::vector<Light> lights = detector_->findLights(src, binary_img);
  std::vector<Armor> armors = detector_->matchLights(lights, enemy_color_);
  std::vector<Armor> real_armors = detector_->numberClassify(armors, src);

  // fill in ros msg
  rm_interfaces::AllTargets targets_msg;
  targets_msg.header = msg->header;
  for (auto &armor : real_armors) {
    rm_interfaces::Armor a;
    a.tl.x = armor.left_light.top.x;
    a.tl.y = armor.left_light.top.y;
    a.bl.x = armor.left_light.bottom.x;
    a.bl.y = armor.left_light.bottom.y;
    a.br.x = armor.right_light.bottom.x;
    a.br.y = armor.right_light.bottom.y;
    a.tr.x = armor.right_light.top.x;
    a.tr.y = armor.right_light.top.y;
    a.id = armor.number;
    a.is_big = armor.is_big;
    targets_msg.armors.emplace_back(a);
  }
  targets_pub_.publish(targets_msg);

  // publish debug info
  if (debug_) {
    ros::Time final_time = ros::Time::now();

    cv::Mat rgb_img_debug = drawResults(src, lights, real_armors);
    double latency = (final_time - start_time).toSec();
    cv::putText(rgb_img_debug, "FPS: " + std::to_string(1.0 / latency),
                cv::Point(10, 30), cv::FONT_HERSHEY_SIMPLEX, 1.0,
                cv::Scalar(0, 255, 0), 2);
    cv::putText(rgb_img_debug,
                "Lightness Thresh: " + std::to_string(lightness_thresh_),
                cv::Point(10, 60), cv::FONT_HERSHEY_SIMPLEX, 1.0,
                cv::Scalar(0, 255, 0), 2);
    rgb_img_pub_.publish(
        cv_bridge::CvImage(msg->header, "bgr8", rgb_img_debug).toImageMsg());

    // concat the number imgs
    if (!real_armors.empty()) {
      std::vector<cv::Mat> number_imgs;
      number_imgs.reserve(real_armors.size());
      for (auto &armor : real_armors) {
        cv::resize(armor.number_img, armor.number_img, cv::Size(28, 28));
        number_imgs.emplace_back(armor.number_img);
      }
      cv::Mat all_num_img;
      cv::vconcat(number_imgs, all_num_img);
      number_img_pub_.publish(
          cv_bridge::CvImage(msg->header, "mono8", all_num_img).toImageMsg());
    }
    binary_img_pub_.publish(
        cv_bridge::CvImage(msg->header, "mono8", binary_img).toImageMsg());
  }
}

void DetectorNode::configure_cb(const detector::DetectorConfig &config,
                                uint32_t level) {
  this->lightness_thresh_ = config.min_lightness;
  this->detector_->set_lightness_thresh(lightness_thresh_);
}

cv::Mat DetectorNode::drawResults(cv::Mat &src,
                                  const std::vector<Light> &lights,
                                  const std::vector<Armor> &armors) {
  cv::Mat image = src.clone();

  for (const auto &light : lights) {
    auto color = (light.color == RED ? cv::Scalar(0, 255, 255)
                                     : cv::Scalar(255, 255, 0));
    cv::ellipse(image, light, color, 2);
  }

  for (const auto &armor : armors) {
    cv::line(image, armor.left_light.top, armor.right_light.bottom,
             cv::Scalar{0, 255, 0}, 1);
    cv::line(image, armor.left_light.bottom, armor.right_light.top,
             cv::Scalar{0, 255, 0}, 1);
    cv::circle(image, armor.left_light.top, 2, cv::Scalar(0, 0, 255), 2);
    cv::circle(image, armor.left_light.bottom, 2, cv::Scalar(0, 255, 0), 2);
    cv::circle(image, armor.right_light.bottom, 2, cv::Scalar(255, 0, 0), 2);
    cv::circle(image, armor.right_light.top, 2, cv::Scalar(0, 255, 255), 2);

    std::string text =
        "id:" + std::to_string(armor.number) + " " +
        std::to_string(static_cast<int>(armor.confidence * 100)) + "%" +
        (armor.is_big ? " BIG" : " SMALL");
    cv::putText(image, text, armor.left_light.top, cv::FONT_HERSHEY_SIMPLEX,
                0.8, cv::Scalar(0, 255, 2));
  }

  return image;
}

}  // namespace detector